{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e7057f38",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/BagelIndexDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f15a2a63",
   "metadata": {},
   "source": [
    "# Bagel Network\n",
    "\n",
    ">[Bagel](https://docs.bageldb.ai/) is a Open Inference Data for AI. It is built for distributed Machine Learning compute. Cutting AI data infra spend by tenfold.\n",
    "\n",
    "<a href=\"https://discord.gg/bA7B6r97\" target=\"_blank\">\n",
    "      <img src=\"https://img.shields.io/discord/1073293645303795742\" alt=\"Discord\">\n",
    "  </a>&nbsp;&nbsp;\n",
    "\n",
    "\n",
    "- [Website](https://www.bageldb.ai/)\n",
    "- [Documentation](https://docs.bageldb.ai/)\n",
    "- [Twitter](https://twitter.com/bageldb_ai)\n",
    "- [Discord](https://discord.gg/bA7B6r97)\n",
    "\n",
    "\n",
    "Install Bagel with:\n",
    "\n",
    "```sh\n",
    "pip install bagelML\n",
    "```\n",
    "\n",
    "\n",
    "Like any other database, you can:\n",
    "- `.add` \n",
    "- `.get` \n",
    "- `.delete`\n",
    "- `.update`\n",
    "- `.upsert`\n",
    "- `.peek`\n",
    "- `.modify`\n",
    "- and `.find` runs the similarity search. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "878a9dc8",
   "metadata": {},
   "source": [
    "## Basic Example\n",
    "\n",
    "In this basic example, we take the a Paul Graham essay, split it into chunks, embed it using an open-source embedding model, load it into Bagel, and then query it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d325a363",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-bagel\n",
    "%pip install llama-index-embeddings-huggingface\n",
    "%pip install bagelML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bb203c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import\n",
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.vector_stores.bagel import BagelVectorStore\n",
    "from llama_index.core import StorageContext\n",
    "from IPython.display import Markdown, display\n",
    "import bagel\n",
    "from bagel import Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bdb122a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up OpenAI\n",
    "import os\n",
    "import getpass\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")\n",
    "import openai\n",
    "\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4456726",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dd2e858",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87cfd091",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create server settings\n",
    "server_settings = Settings(\n",
    "    bagel_api_impl=\"rest\", bagel_server_host=\"api.bageldb.ai\"\n",
    ")\n",
    "\n",
    "# create client\n",
    "client = bagel.Client(server_settings)\n",
    "\n",
    "# create collection\n",
    "collection = client.get_or_create_cluster(\n",
    "    \"testing_embeddings\", embedding_model=\"custom\", dimension=384\n",
    ")\n",
    "\n",
    "# define embedding function\n",
    "embed_model = \"local:BAAI/bge-small-en-v1.5\"\n",
    "\n",
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()\n",
    "\n",
    "# set up BagelVectorStore and load in data\n",
    "vector_store = BagelVectorStore(collection=collection)\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context, embed_model=embed_model\n",
    ")\n",
    "\n",
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "print(f\"<b>{response}</b>\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c45be6f8",
   "metadata": {},
   "source": [
    "## Create - Add - Get"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ffc8240",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_add_get(client):\n",
    "    \"\"\"\n",
    "    Create, add, and get\n",
    "    \"\"\"\n",
    "    name = \"testing\"\n",
    "\n",
    "    # Get or create a cluster\n",
    "    cluster = client.get_or_create_cluster(name)\n",
    "\n",
    "    # Add documents to the cluster\n",
    "    resp = cluster.add(\n",
    "        documents=[\n",
    "            \"This is document1\",\n",
    "            \"This is bidhan\",\n",
    "        ],\n",
    "        metadatas=[{\"source\": \"google\"}, {\"source\": \"notion\"}],\n",
    "        ids=[str(uuid.uuid4()), str(uuid.uuid4())],\n",
    "    )\n",
    "\n",
    "    # Print count\n",
    "    print(\"count of docs:\", cluster.count())\n",
    "\n",
    "    # Get the first item\n",
    "    first_item = cluster.peek(1)\n",
    "    if first_item:\n",
    "        print(\"get 1st item\")\n",
    "\n",
    "    print(\">> create_add_get done !\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0517c3da",
   "metadata": {},
   "source": [
    "## Create - Add - Find by Text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9047a888",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_add_find(client):\n",
    "    \"\"\"\n",
    "    Create, add, & find\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    api : _type_\n",
    "        _description_\n",
    "    \"\"\"\n",
    "    name = \"testing\"\n",
    "\n",
    "    # Get or create a cluster\n",
    "    cluster = client.get_or_create_cluster(name)\n",
    "\n",
    "    # Add documents to the cluster\n",
    "    cluster.add(\n",
    "        documents=[\n",
    "            \"This is document\",\n",
    "            \"This is Towhid\",\n",
    "            \"This is text\",\n",
    "        ],\n",
    "        metadatas=[\n",
    "            {\"source\": \"notion\"},\n",
    "            {\"source\": \"notion\"},\n",
    "            {\"source\": \"google-doc\"},\n",
    "        ],\n",
    "        ids=[str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())],\n",
    "    )\n",
    "\n",
    "    # Query the cluster for similar results\n",
    "    results = cluster.find(\n",
    "        query_texts=[\"This\"],\n",
    "        n_results=5,\n",
    "        where={\"source\": \"notion\"},\n",
    "        where_document={\"$contains\": \"is\"},\n",
    "    )\n",
    "\n",
    "    print(results)\n",
    "    print(\">> create_add_find done  !\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "222e7f85",
   "metadata": {},
   "source": [
    "## Create - Add - Find by Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4de58cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_add_find_em(client):\n",
    "    \"\"\"Create, add, & find embeddings\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    api : _type_\n",
    "        _description_\n",
    "    \"\"\"\n",
    "    name = \"testing_embeddings\"\n",
    "    # Reset the Bagel server\n",
    "    client.reset()\n",
    "\n",
    "    # Get or create a cluster\n",
    "    cluster = api.get_or_create_cluster(name)\n",
    "    # Add embeddings and other data to the cluster\n",
    "    cluster.add(\n",
    "        embeddings=[\n",
    "            [1.1, 2.3, 3.2],\n",
    "            [4.5, 6.9, 4.4],\n",
    "            [1.1, 2.3, 3.2],\n",
    "            [4.5, 6.9, 4.4],\n",
    "            [1.1, 2.3, 3.2],\n",
    "            [4.5, 6.9, 4.4],\n",
    "            [1.1, 2.3, 3.2],\n",
    "            [4.5, 6.9, 4.4],\n",
    "        ],\n",
    "        metadatas=[\n",
    "            {\"uri\": \"img1.png\", \"style\": \"style1\"},\n",
    "            {\"uri\": \"img2.png\", \"style\": \"style2\"},\n",
    "            {\"uri\": \"img3.png\", \"style\": \"style1\"},\n",
    "            {\"uri\": \"img4.png\", \"style\": \"style1\"},\n",
    "            {\"uri\": \"img5.png\", \"style\": \"style1\"},\n",
    "            {\"uri\": \"img6.png\", \"style\": \"style1\"},\n",
    "            {\"uri\": \"img7.png\", \"style\": \"style1\"},\n",
    "            {\"uri\": \"img8.png\", \"style\": \"style1\"},\n",
    "        ],\n",
    "        documents=[\n",
    "            \"doc1\",\n",
    "            \"doc2\",\n",
    "            \"doc3\",\n",
    "            \"doc4\",\n",
    "            \"doc5\",\n",
    "            \"doc6\",\n",
    "            \"doc7\",\n",
    "            \"doc8\",\n",
    "        ],\n",
    "        ids=[\"id1\", \"id2\", \"id3\", \"id4\", \"id5\", \"id6\", \"id7\", \"id8\"],\n",
    "    )\n",
    "\n",
    "    # Query the cluster for results\n",
    "    results = cluster.find(query_embeddings=[[1.1, 2.3, 3.2]], n_results=5)\n",
    "\n",
    "    print(\"find result:\", results)\n",
    "    print(\">> create_add_find_em done  !\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13a08fa4",
   "metadata": {},
   "source": [
    "## Create - Add - Modify - Update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64b0ec14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_add_modify_update(client):\n",
    "    \"\"\"\n",
    "    Create, add, modify, and update\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    api : _type_\n",
    "        _description_\n",
    "    \"\"\"\n",
    "    name = \"testing\"\n",
    "    new_name = \"new_\" + name\n",
    "\n",
    "    # Get or create a cluster\n",
    "    cluster = client.get_or_create_cluster(name)\n",
    "\n",
    "    # Modify the cluster name\n",
    "    print(\"Before:\", cluster.name)\n",
    "    cluster.modify(name=new_name)\n",
    "    print(\"After:\", cluster.name)\n",
    "\n",
    "    # Add documents to the cluster\n",
    "    cluster.add(\n",
    "        documents=[\n",
    "            \"This is document1\",\n",
    "            \"This is bidhan\",\n",
    "        ],\n",
    "        metadatas=[{\"source\": \"notion\"}, {\"source\": \"google\"}],\n",
    "        ids=[\"id1\", \"id2\"],\n",
    "    )\n",
    "\n",
    "    # Retrieve document metadata before updating\n",
    "    print(\"Before update:\")\n",
    "    print(cluster.get(ids=[\"id1\"]))\n",
    "\n",
    "    # Update document metadata\n",
    "    cluster.update(ids=[\"id1\"], metadatas=[{\"source\": \"google\"}])\n",
    "\n",
    "    # Retrieve document metadata after updating\n",
    "    print(\"After update source:\")\n",
    "    print(cluster.get(ids=[\"id1\"]))\n",
    "\n",
    "    print(\">> create_add_modify_update done !\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cbc867a",
   "metadata": {},
   "source": [
    "## Create - Upsert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "587ccb49",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_upsert(client):\n",
    "    \"\"\"\n",
    "    Create and upsert\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    api : _type_\n",
    "        _description_\n",
    "    \"\"\"\n",
    "    # Reset the Bagel server\n",
    "    api.reset()\n",
    "\n",
    "    name = \"testing\"\n",
    "\n",
    "    # Get or create a cluster\n",
    "    cluster = client.get_or_create_cluster(name)\n",
    "\n",
    "    # Add documents to the cluster\n",
    "    cluster.add(\n",
    "        documents=[\n",
    "            \"This is document1\",\n",
    "            \"This is bidhan\",\n",
    "        ],\n",
    "        metadatas=[{\"source\": \"notion\"}, {\"source\": \"google\"}],\n",
    "        ids=[\"id1\", \"id2\"],\n",
    "    )\n",
    "\n",
    "    # Upsert documents in the cluster\n",
    "    cluster.upsert(\n",
    "        documents=[\n",
    "            \"This is document\",\n",
    "            \"This is google\",\n",
    "        ],\n",
    "        metadatas=[{\"source\": \"notion\"}, {\"source\": \"google\"}],\n",
    "        ids=[\"id1\", \"id3\"],\n",
    "    )\n",
    "\n",
    "    # Print the count of documents in the cluster\n",
    "    print(\"Count of documents:\", cluster.count())\n",
    "    print(\">> create_upsert done !\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
